import cv2
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np


def generate_data(sample_num: int, f, width: int, noise: float) -> np.ndarray:
    """
    根据二维的函数产生一个三维的数据分布，其中二维函数在XOZ平面内，Y方向上有随机偏移

    :param sample_num: 样本个数
    :param f: 二维函数z=f(x)
    :param width: 同以x，z坐标下y轴上的点的个数
    :param noise: 噪声比
    :return:x_train:shape=(sample_num,3)
    """
    n = int(sample_num / width)
    x_noise = np.random.normal(0.0, noise, n)
    z_noise = np.random.normal(0.0, noise * 100, n)
    x = np.arange(-5, 5, 10 / n)
    x = x + x_noise
    z = f(x) + z_noise
    y_gap = np.array(range(width))
    y = []
    for i in y_gap:
        y += (np.array([i] * n) + np.random.normal(0.0, noise * 10, n)).tolist()
    x = np.tile(x, width)  # 复制数据
    z = np.tile(z, width)
    y = np.array(y)
    return np.c_[x, y, z]


def draw_picture_3d(x_train: np.ndarray, angle: int, title: str = 'picture'):
    """
    绘制三维散点图

    :param title: 标题名称
    :param angle: 坐标旋转角度
    :param x_train: 三维向量
    :return: 无
    """
    fig = plt.figure()
    ax = Axes3D(fig, auto_add_to_figure=False)
    fig.add_axes(ax)
    x_min = min(x_train[:, 0]) - 1
    x_max = max(x_train[:, 0]) + 1
    y_min = min(x_train[:, 1]) - 1
    y_max = max(x_train[:, 1]) + 1
    z_min = min(x_train[:, 2]) - 1
    z_max = max(x_train[:, 2]) + 1
    # 规定坐标轴的值
    ax.set_xbound(x_min, x_max)
    ax.set_ybound(y_min, y_max)
    ax.set_zbound(z_min, z_max)
    ax.scatter(x_train[:, 0], x_train[:, 1], x_train[:, 2], c='green', marker='o')
    ax.set_title(title)
    ax.set_zlabel('z', fontdict={'size': 15, 'color': 'blue'})
    ax.set_ylabel('y', fontdict={'size': 15, 'color': 'blue'})
    ax.set_xlabel('x', fontdict={'size': 15, 'color': 'blue'})
    # 将图分别上下旋转和左右旋转
    ax.view_init(angle, -90)
    plt.show()


def draw_picture_2d(x_train: np.ndarray, title: str = 'picture'):
    """
    绘制二维散点图

    :param title: 标题名称
    :param x_train: 二维向量
    :return: 无
    """
    x_min = min(x_train[:, 0]) - 1
    x_max = max(x_train[:, 0]) + 1
    y_min = min(x_train[:, 1]) - 1
    y_max = max(x_train[:, 1]) + 1
    plt.axis([x_min, x_max, y_min, y_max])  # 规定坐标轴的值
    plt.scatter(x=x_train[:, 0], y=x_train[:, 1], c='green', marker='o')
    plt.title(title)
    plt.ylabel('y')
    plt.xlabel('x')
    plt.show()


def principal_component_analysis(x_train: np.ndarray, reduced_dimension: int) -> (np.ndarray, np.ndarray, np.ndarray):
    """
    将shape=(sample_num, dimension)的x_train降为shape=(sample_num, reduced_dimension)的新的x_train
    :param x_train: shape=(sample_num, dimension)的需要降维维的数据
    :param reduced_dimension: 需要降到的维度（reduced_dimension<dimension）
    :return: shape=(dimension, reduced_dimension)的特征向量矩阵dimensionality_reduction
    shape=(sample_num, dimension)的中心化的原样本矩阵data_covariance
    shape=(1, dimension)的原样本均值矩阵data_mean
    """
    N = x_train.shape[0]
    data_mean = np.sum(x_train, axis=0) / N  # 计算均值
    data_covariance = x_train - data_mean  # 计算中心化的X值
    covariance_matrix = (data_covariance.T.dot(data_covariance)) / N  # 计算协方差矩阵
    eigenvalue, eigenvector = np.linalg.eig(covariance_matrix)  # 计算特征向量
    index = np.argsort(eigenvalue)  # 取最大的特征值标号
    index = index[::-1]
    dimensionality_reduction = eigenvector[:, index[:reduced_dimension]]  # 取出最大特征值对应的特征向量，拼接成空间变换矩阵
    dimensionality_reduction = np.real(dimensionality_reduction)  # 特征向量可能为复向量，对其保留实部即可
    return dimensionality_reduction, data_covariance, data_mean


def reduce_and_restore(sample_num, width, noise, f):
    """
    展示降维和复原
    :param f: XOZ平面中的函数
    :return: 无
    """
    x_train = generate_data(sample_num, f, width, noise)
    draw_picture_3d(x_train, 45, title='origin')
    draw_picture_3d(x_train, 0, title='origin')
    dimensionality_reduction, data_covariance, data_mean = principal_component_analysis(x_train, reduced_dimension)
    # 降维
    x_train_reduced = data_covariance.dot(dimensionality_reduction)
    draw_picture_2d(x_train_reduced, title='reduced_dimension')
    # 复原
    x_train_restored = x_train_reduced.dot(dimensionality_reduction.T) + data_mean
    draw_picture_3d(x_train_restored, 45, title='restored_dimension')
    draw_picture_3d(x_train_restored, 0, title='restored_dimension')


def picture_demo(file_paths: [], reduced_dimension: int, shape=(40, 40)):
    x_train = []
    for i in range(len(file_paths)):
        plt.subplot(2, 3, i + 1)
        img = cv2.imread(file_paths[i])
        img = cv2.resize(img, shape)  # 压缩图片，提高降维速度
        img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 转换为灰度图
        plt.imshow(img_gray)  # 展示图片
        x_train.append(img_gray.reshape(shape[0] * shape[1]))
    plt.show()
    x_train = np.array(x_train)
    dimensionality_reduction, data_covariance, data_mean = principal_component_analysis(x_train, reduced_dimension)
    # 降维
    x_train_reduced = data_covariance.dot(dimensionality_reduction)
    # 复原
    x_train_restored = x_train_reduced.dot(dimensionality_reduction.T) + data_mean
    for i in range(len(x_train_restored)):
        plt.subplot(2, 3, i + 1)
        img_gray = x_train_restored[i].reshape(shape)
        stnr = signal_to_noise_ratio(x_train[i].reshape(shape), img_gray)
        plt.title('stnr={}db'.format(round(stnr, 3)))
        plt.imshow(img_gray)
    plt.show()


def signal_to_noise_ratio(img1: np.ndarray, img2: np.ndarray) -> float:
    """
    计算灰度图片压缩之后再复原的信噪比

    :param img1: 图片1
    :param img2: 图片2
    :return: 信噪比
    """
    mse = np.mean((img1 / 255. - img2 / 255.) ** 2) / (img1.shape[0] * img1.shape[1])
    return abs(10 * np.math.log10(1 / mse))


if __name__ == '__main__':
    # 初始化参数
    np.random.seed(1)  # 随机数种子
    reduced_dimension = 2

    # # XOZ平面中为一次函数的3d平面
    # sample_num = 2000
    # width = 25
    # noise = 0.005
    # reduce_and_restore(sample_num, width, noise, np.poly1d([3, 5]))

    # XOZ平面中为多项式函数的3d曲面
    # sample_num = 3000
    # width = 50
    # # noise = 0.1
    # noise = 0.001
    # reduce_and_restore(sample_num, width, noise, np.poly1d([1, 0, 0]))

    # 人脸降维并且重建，原图像为40*40=1600维
    file_paths = ['1.jpg', '2.jpg', '3.jpg', '4.jpg', '5.jpg', '6.jpg']
    shape = (40, 40)
    reduced_dimension = 1
    picture_demo(file_paths, reduced_dimension, shape=shape)
